# pip install transformers -i https://pypi.tuna.tsinghua.edu.cn/simple/

from transformers import AutoModel

model = AutoModel.from_pretrained(
    # 'jinaai/jina-reranker-v3',
    "/Users/hhwang/models/jina-reranker-v3/",
    dtype="auto",
    trust_remote_code=True,
)
model.eval()

query = "What are the health benefits of green tea?"
documents = [
    "Green tea contains antioxidants called catechins that may help reduce inflammation and protect cells from damage.",
    "El precio del café ha aumentado un 20% este año debido a problemas en la cadena de suministro.",
    "Studies show that drinking green tea regularly can improve brain function and boost metabolism.",
    "Basketball is one of the most popular sports in the United States.",
    "绿茶富含儿茶素等抗氧化剂，可以降低心脏病风险，还有助于控制体重。",
    "Le thé vert est riche en antioxydants et peut améliorer la fonction cérébrale.",
]

# Rerank documents
results = model.rerank(query, documents)

# Results are sorted by relevance score (highest first)
for result in results:
    print(f"Score: {result['relevance_score']:.4f}")
    print(f"Document: {result['document'][:100]}...")
    print()

# Output:
# Score: 0.2976
# Document: Green tea contains antioxidants called catechins that may help reduce inflammation and protect ce...
#
# Score: 0.2258
# Document: 绿茶富含儿茶素等抗氧化剂，可以降低心脏病风险，还有助于控制体重。
#
# Score: 0.1911
# Document: Studies show that drinking green tea regularly can improve brain function and boost metabolism.
#
# Score: 0.1640
# Document: Le thé vert est riche en antioxydants et peut améliorer la fonction cérébrale.

# curl -X POST https://api.jina.ai/v1/rerank -H "Content-Type: application/json"  -H "Authorization: Bearer JINA_API_KEY" -d '{  "model": "jina-reranker-v3", "query": "slm markdown", "documents": ["slm markdown is something", "Studies show that drinking green tea"], "return_documents": false }'

